"""
训练入口 - 简洁版（类似YOLO）
支持配置文件和命令行参数
"""
import argparse
import yaml
import torch
import time
from pathlib import Path

from configs.default import DefaultConfig
from models import build_model
from core.trainer import Trainer
from utils.general import set_seed, check_requirements, colorstr, select_device


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='Train Meniscus Volume Prediction Model')
    
    # 配置文件
    parser.add_argument('--config', type=str, default=None,
                       help='path to config file (YAML)')
    
    # 模型参数
    parser.add_argument('--model', type=str, default='baseline_cnn',
                       choices=['baseline_cnn', 'deep_cnn', 'gcnet'],
                       help='model type')
    
    # 数据参数
    parser.add_argument('--data', type=str, default='/root/workspace',
                       help='data directory')
    parser.add_argument('--image-size', type=int, nargs=2, default=[288, 96],
                       help='image size [width, height]')
    
    # 训练参数
    parser.add_argument('--epochs', type=int, default=None,
                       help='number of epochs')
    parser.add_argument('--batch-size', type=int, default=None,
                       help='batch size')
    parser.add_argument('--lr', '--learning-rate', type=float, default=None,
                       help='learning rate')
    parser.add_argument('--workers', type=int, default=None,
                       help='number of dataloader workers')
    
    # 优化器和调度器
    parser.add_argument('--optimizer', type=str, default=None,
                       choices=['adam', 'adamw', 'sgd'],
                       help='optimizer type')
    parser.add_argument('--scheduler', type=str, default=None,
                       choices=['plateau', 'cosine', 'step'],
                       help='learning rate scheduler')
    
    # 数据增强
    parser.add_argument('--augment', action='store_true',
                       help='enable data augmentation')
    parser.add_argument('--no-augment', dest='augment', action='store_false',
                       help='disable data augmentation')
    parser.set_defaults(augment=None)
    
    # 设备和其他
    parser.add_argument('--device', type=str, default='',
                       help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--seed', type=int, default=None,
                       help='random seed')
    parser.add_argument('--resume', type=str, default=None,
                       help='resume from checkpoint')
    parser.add_argument('--name', type=str, default=None,
                       help='experiment name')
    parser.add_argument('--save-period', type=int, default=None,
                       help='save checkpoint every x epochs')
    
    return parser.parse_args()


def load_config_from_yaml(yaml_path):
    """
    从YAML文件加载配置
    
    参数:
        yaml_path: YAML配置文件路径
    
    返回:
        config_dict: 配置字典
    """
    with open(yaml_path, 'r', encoding='utf-8') as f:
        config_dict = yaml.safe_load(f)
    return config_dict


def merge_config(args):
    """
    合并配置：默认配置 -> YAML配置 -> 命令行参数
    
    参数:
        args: 命令行参数
    
    返回:
        config: 配置对象
    """
    # 1. 加载默认配置
    config = DefaultConfig()
    
    # 2. 如果有YAML配置文件，加载并更新
    if args.config:
        print(f"\nLoading config from: {args.config}")
        yaml_config = load_config_from_yaml(args.config)
        
        # 更新配置
        for key, value in yaml_config.items():
            if isinstance(value, dict):
                # 处理嵌套字典
                for sub_key, sub_value in value.items():
                    if hasattr(config, sub_key):
                        setattr(config, sub_key, sub_value)
            else:
                if hasattr(config, key):
                    setattr(config, key, value)
    
    # 3. 命令行参数覆盖（最高优先级）
    if args.model:
        config.model_type = args.model
    if args.data:
        config.data_dir = args.data
    if args.image_size:
        config.image_size = tuple(args.image_size)
    if args.epochs is not None:
        config.epochs = args.epochs
    if args.batch_size is not None:
        config.batch_size = args.batch_size
    if args.lr is not None:
        config.learning_rate = args.lr
    if args.workers is not None:
        config.num_workers = args.workers
    if args.optimizer:
        config.optimizer = args.optimizer
    if args.scheduler:
        config.scheduler_type = args.scheduler
    if args.augment is not None:
        config.augment = args.augment
    if args.device:
        config.device = args.device
    if args.seed is not None:
        config.seed = args.seed
    if args.name:
        config.experiment_name = args.name
    if args.save_period is not None:
        config.save_period = args.save_period
    
    return config


def main():
    """主函数"""
    # 解析参数
    args = parse_args()
    
    # 检查依赖
    check_requirements()
    
    # 合并配置
    config = merge_config(args)
    
    # 选择设备
    config.device = select_device(config.device)
    
    # 设置随机种子
    set_seed(config.seed)
    
    # 打印配置
    print(colorstr('bright_blue', 'bold', '\nConfiguration:'))
    config.print_config()
    
    # 构建模型
    print(colorstr('bright_green', '\nBuilding model...'))
    model = build_model(config).to(config.device)
    
    # 打印模型信息
    from utils.general import print_model_info
    print_model_info(model)
    
    # 创建训练器
    trainer = Trainer(model, config)
    
    # 恢复训练
    start_epoch = 0
    if args.resume:
        print(colorstr('bright_yellow', f'\nResuming from checkpoint: {args.resume}'))
        start_epoch = trainer.load_checkpoint(args.resume)
    
    # 开始训练
    print(colorstr('bright_green', 'bold', f'\n{"="*70}'))
    print(colorstr('bright_green', 'bold', 'Starting Training'.center(70)))
    print(colorstr('bright_green', 'bold', f'{"="*70}\n'))
    
    start_time = time.time()
    
    try:
        trainer.train(start_epoch)
    except KeyboardInterrupt:
        print(colorstr('bright_yellow', '\n\nTraining interrupted by user'))
    except Exception as e:
        print(colorstr('bright_red', f'\n\nTraining failed with error: {e}'))
        raise
    
    total_time = time.time() - start_time
    
    # 训练完成
    print(colorstr('bright_green', 'bold', f'\n{"="*70}'))
    print(colorstr('bright_green', 'bold', 'Training Completed!'.center(70)))
    print(colorstr('bright_green', 'bold', f'{"="*70}'))
    print(f'\nTotal training time: {total_time:.2f}s ({total_time/3600:.2f}h)')
    print(f'Results saved to: {config.results_dir}')
    print(colorstr('bright_green', 'bold', f'{"="*70}\n'))


if __name__ == '__main__':
    main()